{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker Clarify E2E Test\n",
    "\n",
    "Simple end-to-end test for the Clarify utils implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[11/08/25 18:39:19] </span><span style=\"color: #0069ff; text-decoration-color: #0069ff; font-weight: bold\">INFO    </span> Found credentials in shared credentials file: ~<span style=\"color: #e100e1; text-decoration-color: #e100e1\">/.aws/credentials</span>   <a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">credentials.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py#1364\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1364</span></a>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[11/08/25 18:39:19]\u001b[0m\u001b[2;36m \u001b[0m\u001b[1;38;2;0;105;255mINFO    \u001b[0m Found credentials in shared credentials file: ~\u001b[38;2;225;0;225m/.aws/\u001b[0m\u001b[38;2;225;0;225mcredentials\u001b[0m   \u001b]8;id=479005;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py\u001b\\\u001b[2mcredentials.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=274417;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py#1364\u001b\\\u001b[2m1364\u001b[0m\u001b]8;;\u001b\\\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import sys\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.datasets import make_classification\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.model_selection import train_test_split\n",
    "import joblib\n",
    "import boto3\n",
    "\n",
    "# Add the clarify utils to path\n",
    "# sys.path.insert(0, '/Users/mollyhe/Documents/SageMaker/sagemaker-python-sdk-staging-molly/sagemaker_utils/src')\n",
    "\n",
    "from sagemaker.core.clarify import (\n",
    "    SageMakerClarifyProcessor,\n",
    "    DataConfig,\n",
    "    BiasConfig,\n",
    "    ModelConfig,\n",
    "    SHAPConfig\n",
    ")\n",
    "from sagemaker.core.helper.session_helper import Session,get_execution_role\n",
    "role = get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Create Sample Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset shape: (1000, 12)\n",
      "Target distribution: target\n",
      "1    503\n",
      "0    497\n",
      "Name: count, dtype: int64\n",
      "Gender distribution: gender\n",
      "0.0    598\n",
      "1.0    402\n",
      "Name: count, dtype: int64\n"
     ]
    }
   ],
   "source": [
    "# Create synthetic dataset\n",
    "X, y = make_classification(\n",
    "    n_samples=1000,\n",
    "    n_features=10,\n",
    "    n_informative=5,\n",
    "    n_redundant=2,\n",
    "    random_state=42\n",
    ")\n",
    "\n",
    "# Add a sensitive feature (simulating gender: 0=female, 1=male)\n",
    "sensitive_feature = np.random.binomial(1, 0.4, size=X.shape[0])\n",
    "X = np.column_stack([X, sensitive_feature])\n",
    "\n",
    "# Create DataFrame\n",
    "feature_names = [f'feature_{i}' for i in range(10)] + ['gender']\n",
    "df = pd.DataFrame(X, columns=feature_names)\n",
    "df['target'] = y\n",
    "\n",
    "print(f\"Dataset shape: {df.shape}\")\n",
    "print(f\"Target distribution: {df['target'].value_counts()}\")\n",
    "print(f\"Gender distribution: {df['gender'].value_counts()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Train Simple Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model accuracy: 0.920\n"
     ]
    }
   ],
   "source": [
    "# Split data\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    df.drop('target', axis=1), df['target'], test_size=0.2, random_state=42\n",
    ")\n",
    "\n",
    "# Train model\n",
    "model = RandomForestClassifier(n_estimators=10, random_state=42)\n",
    "model.fit(X_train, y_train)\n",
    "\n",
    "print(f\"Model accuracy: {model.score(X_test, y_test):.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Upload Data to S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[11/08/25 18:39:30] </span><span style=\"color: #0069ff; text-decoration-color: #0069ff; font-weight: bold\">INFO    </span> Found credentials in shared credentials file: ~<span style=\"color: #e100e1; text-decoration-color: #e100e1\">/.aws/credentials</span>   <a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">credentials.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py#1364\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1364</span></a>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[11/08/25 18:39:30]\u001b[0m\u001b[2;36m \u001b[0m\u001b[1;38;2;0;105;255mINFO    \u001b[0m Found credentials in shared credentials file: ~\u001b[38;2;225;0;225m/.aws/\u001b[0m\u001b[38;2;225;0;225mcredentials\u001b[0m   \u001b]8;id=514072;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py\u001b\\\u001b[2mcredentials.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=337498;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py#1364\u001b\\\u001b[2m1364\u001b[0m\u001b]8;;\u001b\\\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[11/08/25 18:39:31] </span><span style=\"color: #0069ff; text-decoration-color: #0069ff; font-weight: bold\">INFO    </span> Found credentials in shared credentials file: ~<span style=\"color: #e100e1; text-decoration-color: #e100e1\">/.aws/credentials</span>   <a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">credentials.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py#1364\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1364</span></a>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[11/08/25 18:39:31]\u001b[0m\u001b[2;36m \u001b[0m\u001b[1;38;2;0;105;255mINFO    \u001b[0m Found credentials in shared credentials file: ~\u001b[38;2;225;0;225m/.aws/\u001b[0m\u001b[38;2;225;0;225mcredentials\u001b[0m   \u001b]8;id=820880;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py\u001b\\\u001b[2mcredentials.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=366626;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/botocore/credentials.py#1364\u001b\\\u001b[2m1364\u001b[0m\u001b]8;;\u001b\\\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data uploaded to: s3://sagemaker-us-west-2-529088288990/clarify-test/data/test_data.csv\n",
      "Output will be saved to: s3://sagemaker-us-west-2-529088288990/clarify-test/output\n"
     ]
    }
   ],
   "source": [
    "# Setup S3 paths\n",
    "session = Session()\n",
    "bucket = session.default_bucket()\n",
    "prefix = 'clarify-test'\n",
    "\n",
    "# Save test data (without target for inference)\n",
    "test_data = X_test.copy()\n",
    "test_data['target'] = y_test\n",
    "test_data.to_csv('/tmp/test_data.csv', index=False)\n",
    "\n",
    "# Save model\n",
    "joblib.dump(model, '/tmp/model.joblib')\n",
    "\n",
    "# Upload to S3\n",
    "s3_client = boto3.client('s3')\n",
    "s3_client.upload_file('/tmp/test_data.csv', bucket, f'{prefix}/data/test_data.csv')\n",
    "s3_client.upload_file('/tmp/model.joblib', bucket, f'{prefix}/model/model.joblib')\n",
    "\n",
    "data_uri = f's3://{bucket}/{prefix}/data/test_data.csv'\n",
    "output_uri = f's3://{bucket}/{prefix}/output'\n",
    "\n",
    "print(f\"Data uploaded to: {data_uri}\")\n",
    "print(f\"Output will be saved to: {output_uri}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Configure Clarify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Configurations created successfully\n"
     ]
    }
   ],
   "source": [
    "# Data configuration\n",
    "data_config = DataConfig(\n",
    "    s3_data_input_path=data_uri,\n",
    "    s3_output_path=output_uri,\n",
    "    label='target',\n",
    "    headers=list(test_data.columns),\n",
    "    dataset_type='text/csv'\n",
    ")\n",
    "\n",
    "# Bias configuration\n",
    "bias_config = BiasConfig(\n",
    "    label_values_or_threshold=[1],  # Positive class\n",
    "    facet_name='gender',\n",
    "    facet_values_or_threshold=[1]   # Male as sensitive group\n",
    ")\n",
    "\n",
    "# SHAP configuration\n",
    "shap_config = SHAPConfig(\n",
    "    baseline=None,  # Auto-generate baseline\n",
    "    num_samples=10,  # Small number for quick test\n",
    "    agg_method='mean_abs'\n",
    ")\n",
    "\n",
    "print(\"Configurations created successfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Create Clarify Processor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[11/08/25 18:39:35] </span><span style=\"color: #0069ff; text-decoration-color: #0069ff; font-weight: bold\">INFO    </span> Ignoring unnecessary instance type: <span style=\"color: #e100e1; text-decoration-color: #e100e1; font-style: italic\">None</span>.                            <a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/image_uris.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">image_uris.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/image_uris.py#529\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">529</span></a>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[11/08/25 18:39:35]\u001b[0m\u001b[2;36m \u001b[0m\u001b[1;38;2;0;105;255mINFO    \u001b[0m Ignoring unnecessary instance type: \u001b[3;38;2;225;0;225mNone\u001b[0m.                            \u001b]8;id=687743;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/image_uris.py\u001b\\\u001b[2mimage_uris.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=190982;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/image_uris.py#529\u001b\\\u001b[2m529\u001b[0m\u001b]8;;\u001b\\\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Clarify processor created with role: arn:aws:iam::529088288990:role/Admin\n"
     ]
    }
   ],
   "source": [
    "# Create Clarify processor\n",
    "clarify_processor = SageMakerClarifyProcessor(\n",
    "    role=role,\n",
    "    instance_count=1,\n",
    "    instance_type='ml.m5.large',\n",
    "    sagemaker_session=session\n",
    ")\n",
    "\n",
    "print(f\"Clarify processor created with role: {role}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Run Pre-training Bias Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[11/08/25 18:39:36] </span><span style=\"color: #0069ff; text-decoration-color: #0069ff; font-weight: bold\">INFO    </span> Analysis Config: <span style=\"font-weight: bold\">{</span><span style=\"color: #008700; text-decoration-color: #008700\">'dataset_type'</span>: <span style=\"color: #008700; text-decoration-color: #008700\">'text/csv'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'headers'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008700; text-decoration-color: #008700\">'feature_0'</span>, <a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/clarify.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">clarify.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/clarify.py#1992\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">1992</span></a>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         <span style=\"color: #008700; text-decoration-color: #008700\">'feature_1'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'feature_2'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'feature_3'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'feature_4'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'feature_5'</span>,       <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">               </span>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         <span style=\"color: #008700; text-decoration-color: #008700\">'feature_6'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'feature_7'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'feature_8'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'feature_9'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'gender'</span>,          <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">               </span>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         <span style=\"color: #008700; text-decoration-color: #008700\">'target'</span><span style=\"font-weight: bold\">]</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'label'</span>: <span style=\"color: #008700; text-decoration-color: #008700\">'target'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'label_values_or_threshold'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">]</span>,        <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">               </span>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         <span style=\"color: #008700; text-decoration-color: #008700\">'facet'</span>: <span style=\"font-weight: bold\">[{</span><span style=\"color: #008700; text-decoration-color: #008700\">'name_or_index'</span>: <span style=\"color: #008700; text-decoration-color: #008700\">'gender'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'value_or_threshold'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">1</span><span style=\"font-weight: bold\">]}]</span>,     <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">               </span>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         <span style=\"color: #008700; text-decoration-color: #008700\">'methods'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008700; text-decoration-color: #008700\">'report'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008700; text-decoration-color: #008700\">'name'</span>: <span style=\"color: #008700; text-decoration-color: #008700\">'report'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'title'</span>: <span style=\"color: #008700; text-decoration-color: #008700\">'Analysis Report'</span><span style=\"font-weight: bold\">}</span>,  <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">               </span>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         <span style=\"color: #008700; text-decoration-color: #008700\">'pre_training_bias'</span>: <span style=\"font-weight: bold\">{</span><span style=\"color: #008700; text-decoration-color: #008700\">'methods'</span>: <span style=\"font-weight: bold\">[</span><span style=\"color: #008700; text-decoration-color: #008700\">'CI'</span>, <span style=\"color: #008700; text-decoration-color: #008700\">'DPL'</span><span style=\"font-weight: bold\">]}}}</span>                      <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">               </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[11/08/25 18:39:36]\u001b[0m\u001b[2;36m \u001b[0m\u001b[1;38;2;0;105;255mINFO    \u001b[0m Analysis Config: \u001b[1m{\u001b[0m\u001b[38;2;0;135;0m'dataset_type'\u001b[0m: \u001b[38;2;0;135;0m'text/csv'\u001b[0m, \u001b[38;2;0;135;0m'headers'\u001b[0m: \u001b[1m[\u001b[0m\u001b[38;2;0;135;0m'feature_0'\u001b[0m, \u001b]8;id=70589;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/clarify.py\u001b\\\u001b[2mclarify.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=487060;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/clarify.py#1992\u001b\\\u001b[2m1992\u001b[0m\u001b]8;;\u001b\\\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[38;2;0;135;0m'feature_1'\u001b[0m, \u001b[38;2;0;135;0m'feature_2'\u001b[0m, \u001b[38;2;0;135;0m'feature_3'\u001b[0m, \u001b[38;2;0;135;0m'feature_4'\u001b[0m, \u001b[38;2;0;135;0m'feature_5'\u001b[0m,       \u001b[2m               \u001b[0m\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[38;2;0;135;0m'feature_6'\u001b[0m, \u001b[38;2;0;135;0m'feature_7'\u001b[0m, \u001b[38;2;0;135;0m'feature_8'\u001b[0m, \u001b[38;2;0;135;0m'feature_9'\u001b[0m, \u001b[38;2;0;135;0m'gender'\u001b[0m,          \u001b[2m               \u001b[0m\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[38;2;0;135;0m'target'\u001b[0m\u001b[1m]\u001b[0m, \u001b[38;2;0;135;0m'label'\u001b[0m: \u001b[38;2;0;135;0m'target'\u001b[0m, \u001b[38;2;0;135;0m'label_values_or_threshold'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1m]\u001b[0m,        \u001b[2m               \u001b[0m\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[38;2;0;135;0m'facet'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m{\u001b[0m\u001b[38;2;0;135;0m'name_or_index'\u001b[0m: \u001b[38;2;0;135;0m'gender'\u001b[0m, \u001b[38;2;0;135;0m'value_or_threshold'\u001b[0m: \u001b[1m[\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\u001b[1m]\u001b[0m,     \u001b[2m               \u001b[0m\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[38;2;0;135;0m'methods'\u001b[0m: \u001b[1m{\u001b[0m\u001b[38;2;0;135;0m'report'\u001b[0m: \u001b[1m{\u001b[0m\u001b[38;2;0;135;0m'name'\u001b[0m: \u001b[38;2;0;135;0m'report'\u001b[0m, \u001b[38;2;0;135;0m'title'\u001b[0m: \u001b[38;2;0;135;0m'Analysis Report'\u001b[0m\u001b[1m}\u001b[0m,  \u001b[2m               \u001b[0m\n",
       "\u001b[2;36m                    \u001b[0m         \u001b[38;2;0;135;0m'pre_training_bias'\u001b[0m: \u001b[1m{\u001b[0m\u001b[38;2;0;135;0m'methods'\u001b[0m: \u001b[1m[\u001b[0m\u001b[38;2;0;135;0m'CI'\u001b[0m, \u001b[38;2;0;135;0m'DPL'\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m                      \u001b[2m               \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span><span style=\"color: #0069ff; text-decoration-color: #0069ff; font-weight: bold\">INFO    </span> Creating processing-job with name                                    <a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/processing.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">processing.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/processing.py#598\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">598</span></a>\n",
       "<span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">                    </span>         Clarify-Pretraining-Bias-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2025</span>-11-09-02-39-36-699                     <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">                 </span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m                   \u001b[0m\u001b[2;36m \u001b[0m\u001b[1;38;2;0;105;255mINFO    \u001b[0m Creating processing-job with name                                    \u001b]8;id=100415;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/processing.py\u001b\\\u001b[2mprocessing.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=219709;file:///Users/mollyhe/.pyenv/versions/3.12.2/lib/python3.12/site-packages/sagemaker/utils/processing.py#598\u001b\\\u001b[2m598\u001b[0m\u001b]8;;\u001b\\\n",
       "\u001b[2;36m                    \u001b[0m         Clarify-Pretraining-Bias-\u001b[1;36m2025\u001b[0m-11-09-02-39-36-699                     \u001b[2m                 \u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Pre-training bias analysis job submitted successfully\n"
     ]
    }
   ],
   "source": [
    "# Run pre-training bias analysis (no model needed)\n",
    "try:\n",
    "    clarify_processor.run_pre_training_bias(\n",
    "        data_config=data_config,\n",
    "        data_bias_config=bias_config,\n",
    "        methods=['CI', 'DPL'],  # Class Imbalance and Difference in Positive Proportions\n",
    "        wait=False,  # Don't wait for completion in test\n",
    "        logs=False\n",
    "    )\n",
    "    print(\"✅ Pre-training bias analysis job submitted successfully\")\n",
    "except Exception as e:\n",
    "    print(f\"❌ Pre-training bias analysis failed: {str(e)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Status: Completed\n"
     ]
    }
   ],
   "source": [
    "# You can go to SageMaker AI console -> Processing jobs and check the job status\n",
    "# Or you can run the below command\n",
    "# Note that it takes ~5min for the job to be complete\n",
    "\n",
    "response = session.sagemaker_client.describe_processing_job(ProcessingJobName='Clarify-Pretraining-Bias-2025-11-09-02-39-36-699')\n",
    "print(f\"Status: {response['ProcessingJobStatus']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Test Configuration Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Bias analysis config generated successfully\n",
      "Config keys: ['dataset_type', 'headers', 'label', 'label_values_or_threshold', 'facet', 'methods']\n",
      "✅ All required keys present in config\n"
     ]
    }
   ],
   "source": [
    "# Test the internal config generation\n",
    "from sagemaker.core.clarify import _AnalysisConfigGenerator\n",
    "\n",
    "try:\n",
    "    # Generate bias config\n",
    "    bias_analysis_config = _AnalysisConfigGenerator.bias_pre_training(\n",
    "        data_config=data_config,\n",
    "        bias_config=bias_config,\n",
    "        methods=['CI', 'DPL']\n",
    "    )\n",
    "    \n",
    "    print(\"✅ Bias analysis config generated successfully\")\n",
    "    print(f\"Config keys: {list(bias_analysis_config.keys())}\")\n",
    "    \n",
    "    # Validate config structure\n",
    "    required_keys = ['dataset_type', 'label_values_or_threshold', 'facet', 'methods']\n",
    "    missing_keys = [key for key in required_keys if key not in bias_analysis_config]\n",
    "    \n",
    "    if missing_keys:\n",
    "        print(f\"❌ Missing required keys: {missing_keys}\")\n",
    "    else:\n",
    "        print(\"✅ All required keys present in config\")\n",
    "        \n",
    "except Exception as e:\n",
    "    print(f\"❌ Config generation failed: {str(e)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Test Schema Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Schema validation passed\n"
     ]
    }
   ],
   "source": [
    "# Test schema validation\n",
    "from sagemaker.core.clarify import ANALYSIS_CONFIG_SCHEMA_V1_0\n",
    "\n",
    "try:\n",
    "    # Validate the generated config\n",
    "    ANALYSIS_CONFIG_SCHEMA_V1_0.validate(bias_analysis_config)\n",
    "    print(\"✅ Schema validation passed\")\n",
    "except Exception as e:\n",
    "    print(f\"❌ Schema validation failed: {str(e)}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
